"""
限制操作符实现
"""

from typing import Optional, Iterator
import time

from ..executor_base import IteratorExecutor
from ..types import Record, RecordBatch, ExecutionContext

class LimitOperator(IteratorExecutor):
    """限制操作符"""
    
    def __init__(self, executor_id: str, context: ExecutionContext,
                 child: IteratorExecutor, limit: int, offset: int = 0):
        super().__init__(executor_id, context)
        self.child = child
        self.limit = limit
        self.offset = offset
        self.current_position = 0
        self.returned_count = 0
        
        # 添加子执行器
        self.add_child(child)
        
        # 验证参数
        if limit < 0:
            raise ValueError("Limit must be non-negative")
        if offset < 0:
            raise ValueError("Offset must be non-negative")
    
    def open(self):
        """打开限制操作符"""
        with self.lock:
            if self.is_open:
                return
            
            self.is_open = True
            self.current_position = 0
            self.returned_count = 0
            
            # 打开子执行器
            self.child.open()
            
            # 创建限制迭代器
            self.current_iterator = self._create_limit_iterator()
    
    def _create_limit_iterator(self) -> Iterator[Record]:
        """创建限制迭代器"""
        while True:
            # 检查是否已达到限制
            if self.returned_count >= self.limit:
                break
            
            child_batch = self.child.next_batch()
            if not child_batch:
                break
            
            for record in child_batch:
                self.check_timeout()
                
                # 更新处理统计
                self.update_stats(rows_processed=1)
                
                # 检查是否在偏移范围内
                if self.current_position < self.offset:
                    self.current_position += 1
                    continue
                
                # 检查是否已达到限制
                if self.returned_count >= self.limit:
                    break
                
                # 返回记录
                self.returned_count += 1
                self.current_position += 1
                self.update_stats(rows_returned=1)
                yield record
    
    def close(self):
        """关闭限制操作符"""
        with self.lock:
            if not self.is_open:
                return
            
            # 关闭子执行器
            self.child.close()
            
            self.is_open = False
            self.current_iterator = None
    
    def get_estimated_rows(self) -> int:
        """获取估算行数"""
        child_estimated = getattr(self.child, 'get_estimated_rows', lambda: 1000)()
        
        # 考虑偏移和限制
        if child_estimated <= self.offset:
            return 0
        
        available_rows = child_estimated - self.offset
        return min(available_rows, self.limit)

class OffsetOperator(IteratorExecutor):
    """偏移操作符（仅跳过指定数量的记录）"""
    
    def __init__(self, executor_id: str, context: ExecutionContext,
                 child: IteratorExecutor, offset: int):
        super().__init__(executor_id, context)
        self.child = child
        self.offset = offset
        self.skipped_count = 0
        
        # 添加子执行器
        self.add_child(child)
        
        # 验证参数
        if offset < 0:
            raise ValueError("Offset must be non-negative")
    
    def open(self):
        """打开偏移操作符"""
        with self.lock:
            if self.is_open:
                return
            
            self.is_open = True
            self.skipped_count = 0
            
            # 打开子执行器
            self.child.open()
            
            # 创建偏移迭代器
            self.current_iterator = self._create_offset_iterator()
    
    def _create_offset_iterator(self) -> Iterator[Record]:
        """创建偏移迭代器"""
        while True:
            child_batch = self.child.next_batch()
            if not child_batch:
                break
            
            for record in child_batch:
                self.check_timeout()
                
                # 更新处理统计
                self.update_stats(rows_processed=1)
                
                # 检查是否需要跳过
                if self.skipped_count < self.offset:
                    self.skipped_count += 1
                    continue
                
                # 返回记录
                self.update_stats(rows_returned=1)
                yield record
    
    def close(self):
        """关闭偏移操作符"""
        with self.lock:
            if not self.is_open:
                return
            
            # 关闭子执行器
            self.child.close()
            
            self.is_open = False
            self.current_iterator = None

class TopNOperator(IteratorExecutor):
    """TopN操作符，结合排序和限制"""
    
    def __init__(self, executor_id: str, context: ExecutionContext,
                 child: IteratorExecutor, 
                 order_by: list, 
                 limit: int, 
                 offset: int = 0):
        super().__init__(executor_id, context)
        self.child = child
        self.order_by = order_by
        self.limit = limit
        self.offset = offset
        self.all_records = []
        
        # 添加子执行器
        self.add_child(child)
    
    def open(self):
        """打开TopN操作符"""
        with self.lock:
            if self.is_open:
                return
            
            self.is_open = True
            
            # 打开子执行器
            self.child.open()
            
            # 加载所有记录并排序
            self._load_and_sort_records()
            
            # 创建TopN迭代器
            self.current_iterator = self._create_topn_iterator()
    
    def _load_and_sort_records(self):
        """加载并排序记录"""
        # 加载所有记录
        while True:
            child_batch = self.child.next_batch()
            if not child_batch:
                break
            
            self.all_records.extend(child_batch)
            self.update_stats(rows_processed=len(child_batch))
        
        # 排序记录
        self.all_records.sort(key=self._create_sort_key)
    
    def _create_sort_key(self, record: Record) -> tuple:
        """创建排序键"""
        key_values = []
        
        for order_spec in self.order_by:
            column = order_spec.get('column', '')
            direction = order_spec.get('direction', 'ASC').upper()
            
            value = record.get(column)
            
            # 处理None值
            if value is None:
                value = '' if direction == 'ASC' else 'zzz'
            elif isinstance(value, (int, float)):
                if direction == 'DESC':
                    value = -value
            elif isinstance(value, str):
                if direction == 'DESC':
                    # 字符串降序需要特殊处理
                    value = ''.join(chr(255 - ord(c)) for c in value)
            
            key_values.append(value)
        
        return tuple(key_values)
    
    def _create_topn_iterator(self) -> Iterator[Record]:
        """创建TopN迭代器"""
        start_index = self.offset
        end_index = min(start_index + self.limit, len(self.all_records))
        
        for i in range(start_index, end_index):
            self.check_timeout()
            self.update_stats(rows_returned=1)
            yield self.all_records[i]
    
    def close(self):
        """关闭TopN操作符"""
        with self.lock:
            if not self.is_open:
                return
            
            # 关闭子执行器
            self.child.close()
            
            # 清理记录
            self.all_records.clear()
            
            self.is_open = False
            self.current_iterator = None

class PaginationOperator(LimitOperator):
    """分页操作符，扩展限制操作符以支持分页"""
    
    def __init__(self, executor_id: str, context: ExecutionContext,
                 child: IteratorExecutor, 
                 page_size: int, 
                 page_number: int = 1):
        # 计算偏移和限制
        if page_number < 1:
            raise ValueError("Page number must be >= 1")
        if page_size <= 0:
            raise ValueError("Page size must be > 0")
        
        offset = (page_number - 1) * page_size
        super().__init__(executor_id, context, child, page_size, offset)
        
        self.page_size = page_size
        self.page_number = page_number
        self.total_records = 0
        self.total_pages = 0
    
    def open(self):
        """打开分页操作符"""
        # 首先计算总记录数（这需要扫描所有数据）
        self._calculate_total_records()
        
        # 然后执行正常的限制操作
        super().open()
    
    def _calculate_total_records(self):
        """计算总记录数"""
        # 注意：这会导致数据被扫描两次，实际实现中可能需要优化
        temp_child = self.child
        temp_child.open()
        
        count = 0
        while True:
            batch = temp_child.next_batch()
            if not batch:
                break
            count += len(batch)
        
        temp_child.close()
        
        self.total_records = count
        self.total_pages = (count + self.page_size - 1) // self.page_size
    
    def get_pagination_info(self) -> dict:
        """获取分页信息"""
        return {
            'page_number': self.page_number,
            'page_size': self.page_size,
            'total_records': self.total_records,
            'total_pages': self.total_pages,
            'has_next_page': self.page_number < self.total_pages,
            'has_previous_page': self.page_number > 1
        }



